from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import joblib
import numpy as np
import pandas as pd
from typing import Dict, Optional
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

# -------------------------- 1. 加载训练好的模型和组件 --------------------------
try:
    # 加载模型文件（确保模型文件路径正确）
    loaded_data = joblib.load('house_price_prediction_model.pkl')
    model = loaded_data['model']
    poly_transformer = loaded_data['poly_transformer']
    poly_features = loaded_data['poly_features']  # ['floor', 'house_area']
    other_features = loaded_data['other_features']  # 独热编码特征列表
    feature_cols = loaded_data['feature_cols']      # 所有特征列
except Exception as e:
    raise RuntimeError(f"模型加载失败：{str(e)}")

# -------------------------- 2. 定义API请求参数模型 --------------------------
# 基础数值型特征（必填）
class NumericFeatures(BaseModel):
    floor: int = Field(..., ge=1, le=100, description="房屋所在楼的总层数（1-100层）")
    house_area: float = Field(..., ge=10, le=500, description="房屋面积（10-500平米）")

# 独热编码特征（每个类别选一个，默认None，需在API中指定）
class CategoricalFeatures(BaseModel):
    # 地址特征（address_前缀的所有字段，选一个值为1，其余为0）
    address: str = Field(..., description=f"房屋具体地址，可选值：{[col.replace('address_', '') for col in feature_cols if col.startswith('address_')]}")
    # 区域特征（area_前缀的所有字段）
    area: str = Field(..., description=f"房屋所在区域，可选值：{[col.replace('area_', '') for col in feature_cols if col.startswith('area_')]}")
    # 朝向特征（direction_前缀的所有字段）
    direction: str = Field(..., description=f"房屋朝向，可选值：{[col.replace('direction_', '') for col in feature_cols if col.startswith('direction_')]}")
    # 房型特征（house_type_前缀的所有字段）
    house_type: str = Field(..., description=f"房屋房型，可选值：{[col.replace('house_type_', '') for col in feature_cols if col.startswith('house_type_')]}")
    # 周边环境特征（surrounding_前缀的所有字段）
    surrounding: str = Field(..., description=f"周边环境，可选值：{[col.replace('surrounding_', '') for col in feature_cols if col.startswith('surrounding_')]}")

# 合并所有请求参数
class HousePredictionRequest(BaseModel):
    numeric_features: NumericFeatures
    categorical_features: CategoricalFeatures

# -------------------------- 3. 创建FastAPI应用 --------------------------
app = FastAPI(
    title="房屋房价预测API",
    description="基于多元线性回归（多项式特征）的房价预测接口",
    version="1.0.0"
)

# 配置允许跨域的前端地址（若依前端的实际地址）
origins = [
    "http://localhost:82",       # 若依前端地址（无端口时）
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    # 增加预检请求的缓存时间（秒）
    max_age=3600
)

# -------------------------- 4. 辅助函数：构建特征字典 --------------------------
def build_feature_dict(numeric: NumericFeatures, categorical: CategoricalFeatures) -> Dict[str, int | float]:
    """
    将API传入的参数转换为模型需要的特征字典（包含独热编码）
    """
    # 1. 数值型特征
    feature_dict = {
        'floor': numeric.floor,
        'house_area': numeric.house_area
    }

    # 2. 地址特征独热编码（选中的设为1，其余设为0）
    address_cols = [col for col in feature_cols if col.startswith('address_')]
    for col in address_cols:
        feature_dict[col] = 1 if col == f"address_{categorical.address}" else 0

    # 3. 区域特征独热编码
    area_cols = [col for col in feature_cols if col.startswith('area_')]
    for col in area_cols:
        feature_dict[col] = 1 if col == f"area_{categorical.area}" else 0

    # 4. 朝向特征独热编码
    direction_cols = [col for col in feature_cols if col.startswith('direction_')]
    for col in direction_cols:
        feature_dict[col] = 1 if col == f"direction_{categorical.direction}" else 0

    # 5. 房型特征独热编码
    house_type_cols = [col for col in feature_cols if col.startswith('house_type_')]
    for col in house_type_cols:
        feature_dict[col] = 1 if col == f"house_type_{categorical.house_type}" else 0

    # 6. 周边环境特征独热编码
    surrounding_cols = [col for col in feature_cols if col.startswith('surrounding_')]
    for col in surrounding_cols:
        feature_dict[col] = 1 if col == f"surrounding_{categorical.surrounding}" else 0

    # 验证所有特征是否齐全
    missing_features = [col for col in feature_cols if col not in feature_dict]
    if missing_features:
        raise HTTPException(status_code=400, detail=f"缺失特征：{missing_features}")

    return feature_dict

# -------------------------- 5. 预测接口 --------------------------

@app.post("/predict/house_price", response_model=Dict[str, float])
async def predict_house_price(request: HousePredictionRequest):
    """
    房价预测接口
    - 请求体：包含数值特征（floor, house_area）和分类特征（address, area等）
    - 返回：预测的房价（万元）
    """
    # 测试用：直接返回一个数值，排除业务逻辑问题
    # return {"predicted_house_price": 123.45}
    try:
        # 1. 构建特征字典
        feature_dict = build_feature_dict(
            numeric=request.numeric_features,
            categorical=request.categorical_features
        )

        # 2. 特征预处理（与训练时一致）
        # 多项式特征变换（仅数值型）
        poly_data = np.array([[feature_dict[feat] for feat in poly_features]])
        poly_transformed = poly_transformer.transform(poly_data)

        # 提取独热编码特征
        other_data = np.array([[feature_dict[feat] for feat in other_features]])

        # 合并特征
        X_processed = np.hstack([poly_transformed, other_data])

        # 3. 模型预测
        predicted_price = model.predict(X_processed)[0]

        # 4. 返回结果（保留2位小数）
        return {"predicted_house_price": round(predicted_price, 2)}

    except HTTPException as e:
        # 已知错误（参数错误、缺失特征等）
        raise e
    except Exception as e:
        # 未知错误
        raise HTTPException(status_code=500, detail=f"预测失败：{str(e)}")

# -------------------------- 6. 健康检查接口 --------------------------
@app.get("/health", response_model=Dict[str, str])
async def health_check():
    """API健康检查"""
    return {"status": "healthy", "message": "房价预测API正常运行"}

# -------------------------- 7. 查看可选特征接口 --------------------------
@app.get("/features/options", response_model=Dict[str, list])
async def get_feature_options():
    """获取所有分类特征的可选值"""
    return {
        "address_options": [col.replace('address_', '') for col in feature_cols if col.startswith('address_')],
        "area_options": [col.replace('area_', '') for col in feature_cols if col.startswith('area_')],
        "direction_options": [col.replace('direction_', '') for col in feature_cols if col.startswith('direction_')],
        "house_type_options": [col.replace('house_type_', '') for col in feature_cols if col.startswith('house_type_')],
        "surrounding_options": [col.replace('surrounding_', '') for col in feature_cols if col.startswith('surrounding_')]
    }

# -------------------------- 运行API服务 --------------------------
if __name__ == "__main__":
    import uvicorn
    # 本地运行：http://127.0.0.1:8000
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)